package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"log"
	rand2 "math/rand"
	"os"
	"time"
)

var (
	EES    []byte
	AESKEY []byte
)

func loadPublic() {
	s := ""
	if server {
		s = "./private.pem"
	} else {
		s = "./public.pem"
	}
	file, err := os.Open(s)
	if err != nil {
		panic(err)
	}
	info, _ := file.Stat()
	EES = make([]byte, info.Size())
	file.Read(EES)
	file.Close()
}
func loadPrivate() {
	s := ""
	if server {
		s = "./private.pem"
	} else {
		s = "./public.pem"
	}
	file, err := os.Open(s)
	if err != nil {
		panic(err)
	}
	defer file.Close()
	info, _ := file.Stat()
	EES = make([]byte, info.Size())
	file.Read(EES)
}

/*
用rsa加密解密文件
*/
func RSAKeyGen(bits int) error {
	privatekey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		fmt.Println("privatekey create failed")
	}
	derStream := x509.MarshalPKCS1PrivateKey(privatekey)
	block := &pem.Block{
		Type:  "RSA Private key",
		Bytes: derStream,
	}
	fs, err := os.OpenFile("private.pem", os.O_CREATE|os.O_WRONLY, 666)
	if err != nil {
		fmt.Println(err.Error())
		return err
	}
	defer fs.Close()
	fs.Truncate(0)
	fs.Seek(0, 0)
	err = pem.Encode(fs, block)

	// -------
	publicKey := privatekey.PublicKey
	X509PublicKey, err := x509.MarshalPKIXPublicKey(&publicKey)
	if err != nil {
		panic(err)
	}
	fs2, err := os.OpenFile("public.pem", os.O_CREATE|os.O_WRONLY, 666)
	if err != nil {
		fmt.Println(err.Error())
		return err
	}
	defer fs2.Close()
	fs2.Truncate(0)
	fs2.Seek(0, 0)
	publicBlock := pem.Block{Type: "RSA Public Key", Bytes: X509PublicKey}
	err = pem.Encode(fs2, &publicBlock)
	return err
}
func randStr(le int) string {
	s := "0123456hjajsdiqwkmxcAKIJKSDHJQWkelwrjijurit"
	o := ""
	for i := 0; i < le; i++ {
		rand2.Seed(time.Now().UnixNano())
		o += string(s[rand2.Intn(len(s))])
		<-time.After(time.Nanosecond)
	}
	return o
}
func AesKeyGen(bits int) {
	key := randStr(bits)
	fs, err := os.OpenFile("private.pem", os.O_CREATE|os.O_WRONLY, 666)
	if err != nil {
		fmt.Println(err.Error())
	}
	defer fs.Close()
	fs.Truncate(0)
	fs.Seek(0, 0)
	fs.WriteString(key)
	fs.Close()
}
func loadAesKey() {
	fs, err := os.Open("./private.pem")
	if err != nil {
		log.Fatalln(err)
	}
	data, _ := ioutil.ReadAll(fs)
	AESKEY = data
	fs.Close()
}
func padding(src []byte, blockSize int) []byte {
	padNum := blockSize - len(src)%blockSize
	pad := bytes.Repeat([]byte{byte(padNum)}, padNum)
	return append(src, pad...)
}

// 去掉填充数据
func unpadding(src []byte) []byte {
	n := len(src)
	unPadNum := int(src[n-1])
	return src[:n-unPadNum]
}

// 加密
func encryptAES(src []byte) []byte {
	if len(src) == 0 {
		return src
	}
	block, err := aes.NewCipher(AESKEY)
	if err != nil {
		return nil
	}
	src = padding(src, block.BlockSize())
	blockMode := cipher.NewCBCEncrypter(block, AESKEY)
	blockMode.CryptBlocks(src, src)
	return src
}
func decryptAES(src []byte) []byte {
	if len(src) == 0 {
		return src
	}
	block, err := aes.NewCipher(AESKEY)
	if err != nil {
		return nil
	}
	blockMode := cipher.NewCBCDecrypter(block, AESKEY)
	blockMode.CryptBlocks(src, src)
	src = unpadding(src)
	return src
}
func RsaEncrypt(plainText []byte) []byte {
	block, _ := pem.Decode(EES)
	publicKeyInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		panic(err)
	}
	publicKey := publicKeyInterface.(*rsa.PublicKey)
	cipherText, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, plainText)
	if err != nil {
		panic(err)
	}
	return cipherText
}
func RsaDecrypt(cipherText []byte) []byte {
	block, _ := pem.Decode(EES)
	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		panic(err)
	}
	plainText, _ := rsa.DecryptPKCS1v15(rand.Reader, privateKey, cipherText)
	return plainText
}
